import requests
import json

# Ollama 的本地 API 地址
OLLAMA_URL = "http://localhost:11434/api/generate"

# 定义一些工具函数
def get_weather(location: str):
    """模拟获取天气的函数"""
    return f"{location} 的天气是 25°C，晴天。"

def send_email(to: str, subject: str, body: str):
    """模拟发送邮件的函数"""
    return f"已发送邮件给 {to}，主题：{subject}，内容：{body}。"

# 定义可用的工具
TOOLS = {
    "get_weather": {
        "description": "获取指定地点的天气",
        "parameters": {
            "location": {"type": "string", "description": "城市或地区名称"}
        }
    },
    "send_email": {
        "description": "发送电子邮件",
        "parameters": {
            "to": {"type": "string", "description": "收件人邮箱地址"},
            "subject": {"type": "string", "description": "邮件主题"},
            "body": {"type": "string", "description": "邮件内容"}
        }
    }
}

# 提示词，告诉模型如何使用工具
TOOL_PROMPT = """
你可以使用以下工具：
{}
请根据用户的需求，生成一个 JSON 对象来调用工具。例如：
{{"funcalltools": "get_weather", "parameters": {{"location": "北京"}}}}
"""

# 将工具描述添加到提示词中
tool_descriptions = "\n".join([f"- {name}: {info['description']}" for name, info in TOOLS.items()])
TOOL_PROMPT = TOOL_PROMPT.format(tool_descriptions)

def call_tool(tool_name: str, parameters: dict):
    """根据工具名称和参数调用工具"""
    if tool_name == "get_weather":
        return get_weather(**parameters)
    elif tool_name == "send_email":
        return send_email(**parameters)
    else:
        raise ValueError(f"未知工具: {tool_name}")

def ollama_function_call(prompt: str):
    """模拟 Function Calling"""
    # 将工具提示词添加到用户输入中
    full_prompt = TOOL_PROMPT + "\n用户输入：" + prompt

    # 发送请求给 Ollama
    data = {
        "model": "deepseek-r1:7b",
        "prompt": full_prompt,
        "stream": False
    }
    response = requests.post(OLLAMA_URL, json=data)

    if response.status_code == 200:
        # 解析模型的输出
        result = response.json()
        model_output = result["response"]

        try:
            # 尝试解析 JSON 对象
            tool_call = json.loads(model_output.strip())
            tool_name = tool_call["funcalltools"]
            parameters = tool_call["parameters"]

            # 调用工具并返回结果
            tool_result = call_tool(tool_name, parameters)
            return tool_result
        except (json.JSONDecodeError, KeyError):
            # 如果模型没有返回有效的 JSON，直接返回原始输出
            return model_output
    else:
        raise Exception(f"请求失败: {response.status_code}, {response.text}")

# 示例调用
user_input = "获取北京的天气"
result = ollama_function_call(user_input)
print(result)  # 输出：北京 的天气是 25°C，晴天。